import numpy as np
import math
import re
from scipy import interpolate
from scipy.optimize import bisect
from scipy import constants as const
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
#from template import *

class Efit:
    def __init__(self):
        self.nr = None
        self.nz = None


    def read_gfile(self, gfile_name):
        self.gfile_name = gfile_name

        gfile = open(gfile_name, "r")

        #read shot, time, imfit, nr, nz
        line = gfile.readline().split()
        self.title = line[0]
        self.date = line[1]
        self.pound_sign = line[2]
        self.shot = line[3]
        self.time = line[4]
        self.imfit = int(line[5])
        self.nr = int(line[6])
        self.nz = int(line[7])
        print("nr nz = ", self.nr, self.nz)

        #read rdim, zdim, rzero, rleft, zmmid
        line = self.line_to_float(gfile.readline())
        self.rdim = line[0]
        self.zdim = line[1]
        self.rzero = line[2]
        self.rleft = line[3]
        self.zmid = line[4]

        self.zdown = self.zmid-self.zdim/2

        #use this data to derive efit grid
        self.rgefit = np.linspace(self.rleft, self.rleft+self.rdim, self.nr)
        self.zgefit = np.linspace(self.zmid-self.zdim/2, self.zmid+self.zdim/2, self.nz)
        self.drefit = self.rdim / (self.nr - 1)
        self.dzefit = self.zdim / (self.nz - 1)
        self.r, self.z = np.meshgrid(self.rgefit, self.zgefit, indexing='ij')

        self.rmin = self.rleft
        self.rmax = self.rleft + self.rdim
        self.zmin = self.zmid-self.zdim/2
        self.zmax = self.zmid+self.zdim/2


        #read rmaxis, zmaxis, psi_axis, psi_bound, bzero
        line = self.line_to_float(gfile.readline())
        self.rmaxis = line[0]
        self.zmaxis = line[1]
        self.psi_axis = line[2]
        self.psi_bound = line[3]
        self.bzero = line[4]
        print("rzero, Bzero: ", self.rzero, self.bzero)
        print("rmaxis, zmaxis, psi_axis, psi_bound:  ", self.rmaxis, self.zmaxis, self.psi_axis, self.psi_bound)

        self.dpsi = (self.psi_bound - self.psi_axis) /(self.nr - 1)
        self.psi_1d = np.zeros(self.nr)
        for i in range(0, self.nr):
            self.psi_1d[i] = self.psi_axis + self.dpsi * i
        print(self.psi_axis, self.psi_bound)

        #read current, psi_axis, xdum, rmaxis, xdum
        line = self.line_to_float(gfile.readline())
        self.cpasma = line[0]
        self.psi_axis = line[1]
        self.xdum = line[2]
        self.rmaxis = line[3]
        self.xdum = line[4]

        #read zmaxis, xdum, psi_bound, xdum, xdum
        line = self.line_to_float(gfile.readline())
        self.zmaxis = line[0]
        self.xdum = line[1]
        self.psi_bound = line[2]
        self.xdum = line[3]
        self.xdum = line[4]
        print(self.psi_axis, self.psi_bound)

        #read fpol
        self.fpol = np.zeros(self.nr)
        self.line_next = 0
        i_nr = 0
        while i_nr + 5 <= self.nr:
            line = self.line_to_float(gfile.readline())
            for i in range(0, 5):
                self.fpol[i_nr] = line[i]
                #print("fpol: ", self.fpol[i_nr])
                i_nr = i_nr + 1
        if i_nr == self.nr:
            pass
        else:
            line = self.line_to_float(gfile.readline())
            i = 0
            while i_nr < self.nr:
                self.fpol[i_nr] = line[i]
                i = i + 1
                i_nr = i_nr + 1
            self.line_next = i
            #print(self.line_next)


        #read pres
        self.pres = np.zeros(self.nr)
        i_nr = 0
        if self.line_next > 0 and self.line_next < 5:
            for i in range(self.line_next, len(line)):
                self.pres[i_nr] = line[i]
                i_nr = i_nr + 1
        self.line_next = 0

        while i_nr + 5 <= self.nr:
            line = self.line_to_float(gfile.readline())
            for i in range(0, 5):
                self.pres[i_nr] = line[i]
                i_nr = i_nr + 1
        if i_nr == self.nr:
            pass
        else:
            line = self.line_to_float(gfile.readline())
            i = 0
            while i_nr < self.nr:
                self.pres[i_nr] = line[i]
                i = i + 1
                i_nr = i_nr + 1
            self.line_next = i
            #print(self.line_next)


        #read ffprim
        self.ffprim = np.zeros(self.nr)
        i_nr = 0
        if self.line_next > 0 and self.line_next < 5:
            for i in range(self.line_next, len(line)):
                self.ffprim[i_nr] = line[i]
                i_nr = i_nr + 1
        self.line_next = 0

        while i_nr + 5 <= self.nr:
            line = self.line_to_float(gfile.readline())
            for i in range(0, 5):
                self.ffprim[i_nr] = line[i]
                i_nr = i_nr + 1
        if i_nr == self.nr:
            pass
        else:
            line = self.line_to_float(gfile.readline())
            i = 0
            while i_nr < self.nr:
                self.ffprim[i_nr] = line[i]
                i = i + 1
                i_nr = i_nr + 1
            self.line_next = i


        #read pprime
        self.pprime = np.zeros(self.nr)
        i_nr = 0
        if self.line_next > 0 and self.line_next < 5:
            for i in range(self.line_next, len(line)):
                self.pprime[i_nr] = line[i]
                i_nr = i_nr + 1
                
        self.line_next = 0

        while i_nr + 5 <= self.nr:
            line = self.line_to_float(gfile.readline())
            for i in range(0, 5):
                self.pprime[i_nr] = line[i]
                i_nr = i_nr + 1
        if i_nr == self.nr:
            pass
        else:
            line = self.line_to_float(gfile.readline())
            i = 0
            while i_nr < self.nr:
                self.pprime[i_nr] = line[i]
                i = i + 1
                i_nr = i_nr + 1
            self.line_next = i



        self.ffprim = -np.sign(self.cpasma) * self.ffprim
        self.pprime = -np.sign(self.cpasma) * self.pprime


        #read psi
        self.psirz = np.zeros((self.nr, self.nz))
        self.nrz = self.nr * self.nz

        i_nr = 0
        if self.line_next > 0 and self.line_next < 5:
            for i in range(self.line_next, len(line)):
                ir = int(i_nr / self.nz)
                iz = i_nr % self.nz
                self.psirz[ir, iz] = line[i]
                i_nr = i_nr + 1
        self.line_next = 0

        while i_nr + 5 <= self.nrz:
            line = self.line_to_float(gfile.readline())
            for i in range(0, 5):
                ir = int(i_nr / self.nz)
                iz = i_nr % self.nz
                #print(ir, iz, type(ir), type(iz))
                self.psirz[ir, iz] = line[i]
                i_nr = i_nr + 1
        if i_nr == self.nrz:
            pass
        else:
            line = self.line_to_float(gfile.readline())
            i = 0
            while i_nr < self.nrz:
                ir = int(i_nr / self.nz)
                iz = i_nr % self.nz
                self.psirz[ir, iz] = line[i]
                i = i + 1
                i_nr = i_nr + 1
            self.line_next = i
        #print(self.psirz[54, 54])
        #self.psirz_transpose
        self.psirz = self.psirz.transpose()

        #read qpsi
        self.qpsi = np.zeros(self.nr)
        i_nr = 0
        if self.line_next > 0 and self.line_next < 5:
            for i in range(self.line_next, len(line)):
                self.qpsi[i_nr] = line[i]
                i_nr = i_nr + 1
        self.line_next = 0

        while i_nr + 5 <= self.nr:
            line = self.line_to_float(gfile.readline())
            for i in range(0, 5):
                self.qpsi[i_nr] = line[i]
                i_nr = i_nr + 1
        if i_nr == self.nr:
            pass
        else:
            line = self.line_to_float(gfile.readline())
            i = 0
            while i_nr < self.nr:
                self.qpsi[i_nr] = line[i]
                i = i + 1
                i_nr = i_nr + 1
            self.line_next = i
        #print(self.qpsi)


        #read nbbbs, limitr
        line = gfile.readline().split()
        self.nbbbs = int(line[0])
        self.limitr = int(line[1])
        #print(self.nbbbs, self.limitr)

        #read rbbbs, zbbbs
        self.rbbbs = np.zeros(self.nbbbs)
        self.zbbbs = np.zeros(self.nbbbs)
        rzbbbs_temp = np.zeros(self.nbbbs * 2)

        i_nr = 0
        while i_nr + 5 <= self.nbbbs * 2:
            line = self.line_to_float(gfile.readline())
            #print(line)
            for i in range(0, len(line)):
                rzbbbs_temp[i_nr] = line[i]
                i_nr = i_nr + 1
        if i_nr == self.nbbbs * 2:
            pass
        else:
            line = self.line_to_float(gfile.readline())
            i = 0
            while i_nr < self.nbbbs * 2:
                rzbbbs_temp[i_nr] = line[i]
                i = i + 1
                i_nr = i_nr + 1
            self.line_next = i
        for i in range(0, self.nbbbs):
            self.rbbbs[i] = rzbbbs_temp[i*2]
            self.zbbbs[i] = rzbbbs_temp[i*2+1]

        #find X point
        self.rz_x_point = np.array([0.0, 0.0])
        self.rz_x_point[1] = self.zbbbs.min()
        for i in range(0, self.zbbbs.shape[0]):
            if self.zbbbs[i] == self.rz_x_point[1]:
                self.rz_x_point[0] = self.rbbbs[i]
                break

        #read rlim, zlim
        self.rlim = np.zeros(self.limitr)
        self.zlim = np.zeros(self.limitr)
        rzlim_temp = np.zeros(self.limitr * 2)

        i_nr = 0
        while i_nr + 5 <= self.limitr * 2:
            line = self.line_to_float(gfile.readline())
            for i in range(0, len(line)):
                rzlim_temp[i_nr] = line[i]
                i_nr = i_nr + 1
        if i_nr == self.limitr * 2:
            pass
        else:
            line = self.line_to_float(gfile.readline())
            #print(line,i_nr, self.limitr*2)
            i = 0
            while i_nr < self.limitr * 2:
                rzlim_temp[i_nr] = line[i]
                i = i + 1
                i_nr = i_nr + 1
            self.line_next = i

        for i in range(0, self.limitr):
            self.rlim[i] = rzlim_temp[i*2]
            self.zlim[i] = rzlim_temp[i*2+1]

    def write_gfile(self, gfile_name):
        gfile = open(gfile_name, "w")

        #write shot, time, imfit, nr, nz
        gfile.write(self.title + " ")
        gfile.write(self.date + " ")
        gfile.write(self.pound_sign + " ")
        gfile.write(self.shot + " ")
        gfile.write(self.time + " ")
        gfile.write(str(self.imfit) + " ")
        gfile.write(str(self.nr) + " ")
        gfile.write(str(self.nz) + " ")
        gfile.write('\n')

        #write rdim, zdim, rzero, rleft, zmmid
        gfile.write("{: .9e}".format(self.rdim) + " ")
        gfile.write("{: .9e}".format(self.zdim) + " ")
        gfile.write("{: .9e}".format(self.rzero) + " ")
        gfile.write("{: .9e}".format(self.rleft) + " ")
        gfile.write("{: .9e}".format(self.zmid) + " ")
        gfile.write('\n')

        #write rmaxis, zmaxis, psi_axis, psi_bound, bzero
        gfile.write("{: .9e}".format(self.rmaxis) + " ")
        gfile.write("{: .9e}".format(self.zmaxis) + " ")
        gfile.write("{: .9e}".format(self.psi_axis) + " ")
        gfile.write("{: .9e}".format(self.psi_bound) + " ")
        gfile.write("{: .9e}".format(self.bzero) + " ")
        gfile.write('\n')

        #write current, psi_axis, xdum, rmaxis, xdum
        gfile.write("{: .9e}".format(self.cpasma) + " ")
        gfile.write("{: .9e}".format(self.psi_axis) + " ")
        gfile.write("{: .9e}".format(self.xdum) + " ")
        gfile.write("{: .9e}".format(self.rmaxis) + " ")
        gfile.write("{: .9e}".format(self.xdum) + " ")
        gfile.write('\n')

        #write zmaxis, xdum, psi_bound, xdum, xdum
        gfile.write("{: .9e}".format(self.zmaxis) + " ")
        gfile.write("{: .9e}".format(self.xdum) + " ")
        gfile.write("{: .9e}".format(self.psi_bound) + " ")
        gfile.write("{: .9e}".format(self.xdum) + " ")
        gfile.write("{: .9e}".format(self.xdum) + " ")
        gfile.write('\n')

        #write fpol
        for i in range(0, self.nr):
            gfile.write("{: .9e}".format(self.fpol[i]) + " ")
        gfile.write('\n')

        #write pres
        for i in range(0, self.nr):
            gfile.write("{: .9e}".format(self.pres[i]) + " ")
        gfile.write('\n')

        #write ffprim
        for i in range(0, self.nr):
            gfile.write("{: .9e}".format(self.ffprim[i]) + " ")
        gfile.write('\n')

        #write pprime
        for i in range(0, self.nr):
            gfile.write("{: .9e}".format(self.pprime[i]) + " ")
        gfile.write('\n')

        #write psirz
        for i in range(0, self.nr):
            for j in range(0, self.nz):
                gfile.write("{: .9e}".format(self.psirz.transpose()[i, j]) + " ")
        gfile.write('\n')

        #write qpsi
        self.qpsi = np.zeros(self.nr)
        for i in range(0, self.nr):
            gfile.write("{: .9e}".format(self.qpsi[i]) + " ")
        gfile.write('\n')

        
        #=========================== divertor =================================
        #write nbbbs, limitr(for divertor)
        gfile.write(str(self.nbbbs) + " ")
        gfile.write(str(self.wall_all_number) + " ")
        gfile.write('\n')

        #write rbbbs, zbbbs
        for i in range(0, self.nbbbs):
            gfile.write("{: .9e}".format(self.rbbbs[i]) + " ")
            gfile.write("{: .9e}".format(self.zbbbs[i]) + " ")
        gfile.write('\n')

        #write rlim, zlim for divertor
        for i in range(0, self.wall_all_number):
            gfile.write("{: .9e}".format(self.wall_all[i, 0]) + " ")
            gfile.write("{: .9e}".format(self.wall_all[i, 1]) + " ")
        gfile.write('\n')
        
        '''
        #=========================== limiter =================================
        #write nbbbs, limitr(for divertor)
        gfile.write(str(self.nbbbs) + " ")
        gfile.write(str(self.limitr) + " ")
        gfile.write('\n')

        #write rbbbs, zbbbs
        for i in range(0, self.nbbbs):
            gfile.write("{: .9e}".format(self.rbbbs[i]) + " ")
            gfile.write("{: .9e}".format(self.zbbbs[i]) + " ")
        gfile.write('\n')

        #write rlim, zlim for divertor
        for i in range(0, self.limitr):
            gfile.write("{: .9e}".format(self.rlim[i]) + " ")
            gfile.write("{: .9e}".format(self.zlim[i]) + " ")
        gfile.write('\n')
        '''


    def line_to_float(self, line_str):
        length = 16
        str_arr = re.findall('.{'+str(length)+'}', line_str)
        str_arr.append(line_str[(len(str_arr)*length):])
        line = []
        for i in range(0, 5):
            try:
                #print("str_arr ",i, " ", str_arr[i])
                line.append(float(str_arr[i]))
            except:
                print("this line only has ", i-1, " elements ")
                break
        if len(line) == 0:
            line = self.split_line_to_float(line_str)
        return line
               
    def split_line_to_float(self,line_str):
        line = []
        str_arr = line_str.split()
        for i in range(0, len(str_arr)):
            line.append(float(str_arr[i]))
        return line

    def read_firstwall_divertor(self, firstwall_inner, firstwall_outer, divertor):
        self.firstwall_inner = np.loadtxt(firstwall_inner)
        self.firstwall_outer = np.loadtxt(firstwall_outer)
        self.divertor = np.loadtxt(divertor)

        #data is complete
        #self.divertor_upper_left = self.divertor[0:24, :].copy()
        #self.divertor_upper_middle = self.divertor[24:52, :].copy()
        #self.divertor_upper_right = self.divertor[52:self.divertor.shape[0], :].copy()

        #some data is cut
        self.divertor_upper_left = self.divertor[1:24, :].copy()
        self.divertor_upper_middle = self.divertor[24:52, :].copy()
        self.divertor_upper_right = self.divertor[53:self.divertor.shape[0]-2, :].copy()
        print("shape of divertor: ",self.divertor.shape, self.divertor_upper_right.shape)
        for i in range(0, 10):
            self.divertor_upper_right[i, 0] = self.divertor[53+i, 0]
            self.divertor_upper_right[i, 1] = self.divertor[53+i, 1]
        for i in range(10, self.divertor_upper_right.shape[0]):
            self.divertor_upper_right[i, 0] = self.divertor[53+i+2, 0]
            self.divertor_upper_right[i, 1] = self.divertor[53+i+2, 1]

        self.divertor_lower_left = self.divertor_upper_left.copy()
        self.divertor_lower_middle = self.divertor_upper_middle.copy()
        self.divertor_lower_right = self.divertor_upper_right.copy()
        for i in range(0, self.divertor_lower_left.shape[0]):
            self.divertor_lower_left[i, 1] = -self.divertor_lower_left[i, 1]
        for i in range(0, self.divertor_lower_middle.shape[0]):
            self.divertor_lower_middle[i, 1] = -self.divertor_lower_middle[i, 1]
        for i in range(0, self.divertor_lower_right.shape[0]):
            self.divertor_lower_right[i, 1] = -self.divertor_lower_right[i, 1]

        number_wall_all = 0
        number_wall_all = number_wall_all + self.divertor_upper_left.shape[0]
        number_wall_all = number_wall_all + self.divertor_upper_middle.shape[0]
        number_wall_all = number_wall_all + self.divertor_upper_right.shape[0]
        number_wall_all = number_wall_all + self.divertor_lower_left.shape[0]
        number_wall_all = number_wall_all + self.divertor_lower_middle.shape[0]
        number_wall_all = number_wall_all + self.divertor_lower_right.shape[0]
        number_wall_all = number_wall_all + self.firstwall_inner.shape[0]
        number_wall_all = number_wall_all + self.firstwall_outer.shape[0]

        self.wall_all = np.zeros((number_wall_all+1, 2))
        self.wall_all_number = number_wall_all + 1
        i = 0

        j = 0
        while j < self.divertor_upper_left.shape[0]:
            self.wall_all[i, 0] = self.divertor_upper_left[j, 0]
            self.wall_all[i, 1] = self.divertor_upper_left[j, 1]
            j = j + 1
            i = i + 1
        
        j = 0
        while j < self.divertor_upper_middle.shape[0]:
            self.wall_all[i, 0] = self.divertor_upper_middle[j, 0]
            self.wall_all[i, 1] = self.divertor_upper_middle[j, 1]
            j = j + 1
            i = i + 1           

        j = 0
        while j < self.divertor_upper_right.shape[0]:
            self.wall_all[i, 0] = self.divertor_upper_right[-1 - j, 0]
            self.wall_all[i, 1] = self.divertor_upper_right[-1 - j, 1]
            j = j + 1
            i = i + 1

        j = 0
        while j < self.firstwall_outer.shape[0]:
            self.wall_all[i, 0] = self.firstwall_outer[-1 - j, 0]
            self.wall_all[i, 1] = self.firstwall_outer[-1 - j, 1]
            j = j + 1
            i = i + 1

        j = 0
        while j < self.divertor_lower_right.shape[0]:
            self.wall_all[i, 0] = self.divertor_lower_right[j, 0]
            self.wall_all[i, 1] = self.divertor_lower_right[j, 1]
            j = j + 1
            i = i + 1
        
        j = 0
        while j < self.divertor_lower_middle.shape[0]:
            self.wall_all[i, 0] = self.divertor_lower_middle[-1 - j, 0]
            self.wall_all[i, 1] = self.divertor_lower_middle[-1 - j, 1]
            j = j + 1
            i = i + 1  

        j = 0
        while j < self.divertor_lower_left.shape[0]:
            self.wall_all[i, 0] = self.divertor_lower_left[-1 - j, 0]
            self.wall_all[i, 1] = self.divertor_lower_left[-1 - j, 1]
            j = j + 1
            i = i + 1

        j = 0
        while j < self.firstwall_inner.shape[0]:
            self.wall_all[i, 0] = self.firstwall_inner[-1 - j, 0]
            self.wall_all[i, 1] = self.firstwall_inner[-1 - j, 1]
            j = j + 1
            i = i + 1
        self.wall_all[-1, 0] = self.wall_all[0, 0]
        self.wall_all[-1, 1] = self.wall_all[0, 1]


        np.savetxt("efit/hl2a_wall_all.txt", self.wall_all)


        '''
        self.wall_all_solps = np.loadtxt("structure_hl2a_solps.ogr")
        for i in range(0, self.wall_all_solps.shape[0]):
            self.wall_all_solps[i, 0] = self.wall_all_solps[i, 0] / 1000.0
            self.wall_all_solps[i, 1] = self.wall_all_solps[i, 1] / 1000.0
        '''
        
        
        #fig=plt.figure(figsize=(5,10))
        #fig.subplots_adjust(top=0.9,bottom=0.1,wspace=0.5,hspace=0.55)

        #ax0=fig.add_subplot(1,1,1)

        '''
        #line0 = ax0.plot(self.rlim, self.zlim, color="blue", linewidth=3)
        line0 = ax0.plot(self.rbbbs, self.zbbbs, color="red", linewidth=2)
        print("self.rbbbs.shape: ", self.rbbbs.shape)

        ax0.scatter(self.rmaxis, self.zmaxis, color = "blue")
        ax0.scatter(self.rbbbs[0], self.zbbbs[0], color = "blue")
        ax0.scatter(self.rbbbs[33], self.zbbbs[33], color = "blue")
        ax0.scatter(self.rzero, 0.0, color = "blue")
        print("lcfs rz at index 32: ", self.rbbbs[32], self.zbbbs[32])
        print("lcfs rz at index 33: ", self.rbbbs[33], self.zbbbs[33])

        #plot firstwall and divertor
        #line0 = ax0.plot(self.firstwall_inner[:, 0], self.firstwall_inner[:, 1], color="black", linewidth=3)
        #line0 = ax0.plot(self.firstwall_outer[:, 0], self.firstwall_outer[:, 1], color="black", linewidth=3)
        #line0 = ax0.plot(self.divertor[:, 0], self.divertor[:, 1], color="black", linewidth=3)
        #line0 = ax0.plot(self.divertor[:, 0], -self.divertor[:, 1], color="black", linewidth=3)
        '''

        #new plot one by one 
        #line0 = ax0.plot(self.divertor_upper_left[:, 0], self.divertor_upper_left[:, 1], color="black", linewidth=3)
        #line0 = ax0.plot(self.divertor_upper_middle[:, 0], self.divertor_upper_middle[:, 1], color="black", linewidth=3)
        #line0 = ax0.plot(self.divertor_upper_right[:, 0], self.divertor_upper_right[:, 1], color="black", linewidth=3)

        #line0 = ax0.plot(self.divertor_lower_left[:, 0], self.divertor_lower_left[:, 1], color="black", linewidth=3)
        #line0 = ax0.plot(self.divertor_lower_middle[:, 0], self.divertor_lower_middle[:, 1], color="black", linewidth=3)
        #line0 = ax0.plot(self.divertor_lower_right[:, 0], self.divertor_lower_right[:, 1], color="black", linewidth=3)

        #new plot all
        #line0 = ax0.plot(self.wall_all[:, 0], self.wall_all[:, 1], color="black", linewidth=3)

        #ax0.set_aspect('equal')
        
        #============== get inner edge ===================
        dr = 0.05
        step = 0.05
        point0 = np.array([self.rz_out_midplane_lcfs[0] - dr, self.rz_out_midplane_lcfs[1]])
        self.inner_edge = self.mflt_2d_inner(point0, step)
        

    def plot_divertor(self):
        fig=plt.figure(figsize=(5,10))
        fig.subplots_adjust(top=0.9,bottom=0.1,wspace=0.5,hspace=0.55)

        ax0=fig.add_subplot(1,1,1)

        #wall with divertor, anticlockwise 

        #first wall inner(left)
        #line0 = ax0.plot(self.firstwall_inner[:, 0], self.firstwall_inner[:, 1], color="black", linewidth=3)
        #line0 = ax0.plot(self.divertor_lower_left[:, 0], self.divertor_lower_left[:, 1], color="black", linewidth=3)
        self.finalwall_fw_inner = np.append(self.firstwall_inner.copy(), self.divertor_lower_left.copy(), axis=0)
        self.finalwall_fw_inner = np.append(self.divertor_upper_left.copy()[::-1], self.finalwall_fw_inner, axis=0)
        self.finalwall_fw_inner[:, 0] = self.finalwall_fw_inner[:, 0] - 0.02
        line0 = ax0.plot(self.finalwall_fw_inner[:, 0], self.finalwall_fw_inner[:, 1], color="black", linewidth=3)
        self.delete_duplicate(self.finalwall_fw_inner)

        #divertor target lower left
        #line0 = ax0.plot(self.divertor_lower_middle[0:3, 0], self.divertor_lower_middle[0:3, 1], color="black", linewidth=3)
        #print(self.divertor_lower_middle[0:3, :])
        self.finalwall_div_ll = self.divertor_lower_middle[1:3, :].copy()
        line0 = ax0.plot(self.finalwall_div_ll[:, 0], self.finalwall_div_ll[:, 1], color="black", linewidth=3)
        self.delete_duplicate(self.finalwall_div_ll)

        #dome lower
        self.finalwall_dome_l = self.divertor_lower_middle[3:-3, :].copy()
        line0 = ax0.plot(self.finalwall_dome_l[:, 0], self.finalwall_dome_l[:, 1], color="black", linewidth=3)
        self.delete_duplicate(self.finalwall_dome_l)

        #divertor target lower right
        self.finalwall_div_lr = self.divertor_lower_middle[-3:-1, :].copy()
        line0 = ax0.plot(self.finalwall_div_lr[:, 0], self.finalwall_div_lr[:, 1], color="black", linewidth=3)
        self.delete_duplicate(self.finalwall_div_lr)

        #first wall outer(right)
        self.finalwall_fw_outer = np.append(self.divertor_lower_right.copy()[::-1], self.firstwall_outer, axis=0)
        self.finalwall_fw_outer = np.append(self.finalwall_fw_outer, self.divertor_upper_right, axis=0)
        self.finalwall_fw_outer[:, 0] = self.finalwall_fw_outer[:, 0] + 0.02
        line0 = ax0.plot(self.finalwall_fw_outer[:, 0], self.finalwall_fw_outer[:, 1], color="black", linewidth=3)
        self.delete_duplicate(self.finalwall_fw_outer)

        #line0 = ax0.plot(self.divertor_upper_middle[:, 0], self.divertor_upper_middle[:, 1], color="black", linewidth=3)
        divertor_upper_middle = self.divertor_upper_middle.copy()[::-1]

        #divertor target upper right
        self.finalwall_div_ur = divertor_upper_middle[1:3, :].copy()
        #print("self.finalwall_div_ur ", self.finalwall_div_ur)
        line0 = ax0.plot(self.finalwall_div_ur[:, 0], self.finalwall_div_ur[:, 1], color="black", linewidth=3)
        self.delete_duplicate(self.finalwall_div_ur)

        #dome upper
        self.finalwall_dome_u = divertor_upper_middle[3:-3, :].copy()
        line0 = ax0.plot(self.finalwall_dome_u[:, 0], self.finalwall_dome_u[:, 1], color="black", linewidth=3)
        self.delete_duplicate(self.finalwall_dome_u)

        #divertor target upper left
        self.finalwall_div_ul = divertor_upper_middle[-3:-1, :].copy()
        print("self.finalwall_div_ul ", self.finalwall_div_ul)
        line0 = ax0.plot(self.finalwall_div_ul[:, 0], self.finalwall_div_ul[:, 1], color="black", linewidth=3)
        self.delete_duplicate(self.finalwall_div_ul)
        
    

        ax0.set_aspect('equal')
        plt.savefig("divertor0.png")
        plt.show()

    def delete_duplicate(self, finalwall_array):
        for i in range(1, finalwall_array.shape[0]):
            if(finalwall_array[i, 0] == finalwall_array[i-1, 0] and finalwall_array[i, 1] == finalwall_array[i-1, 1]):
                print("find duplicate coord: ", finalwall_array[i, :])

    def init_interp(self):
        self.psirz_transpose = self.psirz.transpose()
        val = np.zeros((self.nr, self.nz))
        for i in range(0, self.nr):
            for j in range(0, self.nz):
                val[i,j] = self.rgefit[i] * self.zgefit[j]
        self.psi_rz_interp2d = interpolate.RectBivariateSpline(self.rgefit, self.zgefit, self.psirz)
        psi0 = self.psi_rz_interp2d(1.6, 0.0)
        psi0_derive = self.psi_rz_interp2d(1.6, 0.0, grid = False, dx = 1)
        psi1 = self.psirz[32, 32]
        psi2 = self.psirz[35, 35]
        
        #print("inter2d test 32 32: ", self.r.shape, self.z.shape, self.psirz.shape, psi0, psi1, psi2, psi0_derive)

        self.f_rz_interp1d = interpolate.interp1d(self.rgefit, self.fpol)

        a = np.array([2.175, 0.0])
        b = self.psi_rz_interp2d(2.175, 0.0, grid = False)
        print(self.r[60,32], self.z[60,32], a, self.psirz[60, 32], b)


        self.Br = np.zeros((self.nr, self.nz))
        self.Bz = np.zeros((self.nr, self.nz))
        self.Bphi = np.zeros((self.nr, self.nz))
        self.Bmag = np.zeros((self.nr, self.nz))
        for i in range(0, self.nr):
            for j in range(0, self.nz):
                r = self.rgefit[i]
                z = self.zgefit[j]
                rz = np.array([r, z])
                Brz = self.get_B_rz_3d(rz)
                self.Br[i, j] = Brz[0]
                self.Bz[i, j] = Brz[1]
                self.Bphi[i, j] = Brz[2]
                self.Bmag[i, j] = math.sqrt(Brz[0] * Brz[0] + Brz[1] * Brz[1] + Brz[2] * Brz[2])

        rz = np.array([self.rmaxis, self.zmaxis])
        Bmaxis = self.get_B_rz_3d(rz)
        print("B at magnetic axis is: ", Bmaxis)

        rz = np.array([self.rzero, 0.0])
        bzeros_interp = self.get_B_rz_3d(rz)
        print("B zero from interp: ", bzeros_interp)

        print("B at magnetic axis from bzerosis : ", self.bzero * self.rzero / self.rmaxis)
        
        self.rz_out_midplane_lcfs = self.get_rz_of_psi(self.psi_bound, 0.0, 1.0)
        print("rz_out_midplane_lcfs", self.rz_out_midplane_lcfs)

 
    #magnetic field line tracizng 2d
    def mflt_2d(self, rz):
        r = []
        z = []
        r.append(rz[0])
        z.append(rz[1])
        dl = -0.002
        for i in range(0, 1000):
            rz0 = np.array([r[i], z[i]])
            rz1 = self.mflt_rk4(rz0, dl, self.mflt_rk4_f)
            r.append(rz1[0])
            z.append(rz1[1])

        mflt_rz = np.array([r, z])
        mflt_rz = mflt_rz.transpose()
        return mflt_rz

    #magnetic field line tracizng 2d in sol, magnetic field line intersect with wall
    def mflt_2d_intersect_wall(self, rz, reverse = False):
        r = []
        z = []
        r.append(rz[0])
        z.append(rz[1])

        dl = 0.002
        if reverse:
            dl = -dl
        is_stop = False
        for i in range(0, 2000):
            rz0 = np.array([r[i], z[i]])
            rz1 = self.mflt_rk4(rz0, dl, self.mflt_rk4_f)
            for j in range(0, self.wall_all.shape[0]-1):
                start_point = np.array([self.wall_all[j, 0], self.wall_all[j, 1]])
                end_point   = np.array([self.wall_all[j+1, 0], self.wall_all[j+1, 1]])
                is_intersect, intersect_point = self.segment_intersect(rz0, rz1, start_point, end_point)
                if is_intersect:
                    is_stop = True
                    break
            if is_stop:
                r.append(intersect_point[0])
                z.append(intersect_point[1])
                break
            else:
                r.append(rz1[0])
                z.append(rz1[1])

        mflt_rz = np.array([r, z])
        mflt_rz = mflt_rz.transpose()
        return mflt_rz


    #magnetic field line tracizng 2d in sol, magnetic field line is closed inside LCFS
    #default line step is 0.001 m
    def mflt_2d_inner(self, rz, step = 0.001, reverse = False):
        r = []
        z = []
        r.append(rz[0])
        z.append(rz[1])

        dl = step
        if reverse:
            dl = -dl
        is_stop = False

        n_cross_lcfs = 0
        for i in range(0, 2000):
            rz0 = np.array([r[i], z[i]])
            rz1 = self.mflt_rk4(rz0, dl, self.mflt_rk4_f)

            if n_cross_lcfs == 0 and rz1[1] < self.rz_out_midplane_lcfs[1]:
                n_cross_lcfs = n_cross_lcfs + 1
            if n_cross_lcfs == 1 and rz1[1] > self.rz_out_midplane_lcfs[1]:
                break

            r.append(rz1[0])
            z.append(rz1[1])


        mflt_rz = np.array([r, z])
        mflt_rz = mflt_rz.transpose()
        return mflt_rz


    def mflt_rk4(self, y0, dt, f):
        k1 = f(y0)
        k2 = f(y0 + 0.5 * dt * k1)
        k3 = f(y0 + 0.5 * dt * k2)
        k4 = f(y0 + dt * k3)
        y1 = y0 + dt * (k1 + 2.0 * k2 + 2.0 * k3 + k4) / 6.0
        #print(y1)
        return y1
    
    def mflt_rk4_f(self, y):
        f = self.get_B_rz(y)
        #print(y, f)
        B_rz_magnitude = math.sqrt(f[0] * f[0] + f[1] * f[1])
        f = f / B_rz_magnitude
        return f

    # B is in cylinder coordinate, r, z, phi
    def get_B_rz(self, rz):
        B = np.zeros(2)
        dpsi_dr = self.psi_rz_interp2d(rz[0], rz[1], dx = 1, grid = False)
        dpsi_dz = self.psi_rz_interp2d(rz[0], rz[1], dy = 1, grid = False)

        B[0] = -dpsi_dz / rz[0]
        B[1] = dpsi_dr / rz[0]
        
        #B[2] = 
        #print(B)
        return B

    #magnetic field line tracizng 3d
    def mflt_3d(self, rz):
        r = []
        z = []
        phi = []
        r.append(rz[0])
        z.append(rz[1])
        phi.append(0.0)
        dl = 0.01
        for i in range(0, 10000):
            rz0 = np.array([r[i], z[i], phi[i]])
            rz1 = self.mflt_rk4(rz0, dl, self.mflt_rk4_f_3d)
            r.append(rz1[0])
            z.append(rz1[1])
            phi.append(rz1[2])

        mflt_rz = np.array([r, z, phi])
        mflt_rz = mflt_rz.transpose()
        for i in range(0, mflt_rz.shape[0]):
            xyz = self.rzphi_to_xyz(mflt_rz[i, :])
            mflt_rz[i, 0] = xyz[0]
            mflt_rz[i, 1] = xyz[1]
            mflt_rz[i, 2] = xyz[2]

        return mflt_rz

    def mflt_rk4_f_3d(self, y):
        f = self.get_B_rz_3d(y)
        #print(y, f)
        B_rz_magnitude = math.sqrt(f[0] * f[0] + f[1] * f[1] + f[2] * f[2])
        f = f / B_rz_magnitude
        return f

    # B is in cylinder coordinate, r, z, phi
    def get_B_rz_3d(self, rz):
        B = np.zeros(3)
        dpsi_dr = self.psi_rz_interp2d(rz[0], rz[1], dx = 1, grid = False)
        dpsi_dz = self.psi_rz_interp2d(rz[0], rz[1], dy = 1, grid = False)
        f = self.f_rz_interp1d(rz[0])

        B[0] = -dpsi_dz / rz[0]
        B[1] = dpsi_dr / rz[0]
        B[2] = f / rz[0]
        
        #B[2] = 
        #print(B)
        return B

    def get_B_unit_rz_3d(self, rz):
        B = np.zeros(3)
        dpsi_dr = self.psi_rz_interp2d(rz[0], rz[1], dx = 1, grid = False)
        dpsi_dz = self.psi_rz_interp2d(rz[0], rz[1], dy = 1, grid = False)
        f = self.f_rz_interp1d(rz[0])

        B[0] = -dpsi_dz / rz[0]
        B[1] = dpsi_dr / rz[0]
        B[2] = f / rz[0]
        
        #B[2] = 
        #print(B)
        return B / np.linalg.norm(B)

    #get psi
    def get_psi(self, rz):
        return self.psi_rz_interp2d(rz[0], rz[1])

    def rzphi_to_xyz(self, rzphi):
        xyz = np.zeros(3)
        xyz[0] = rzphi[0] * math.cos(rzphi[2])
        xyz[1] = rzphi[0] * math.sin(rzphi[2])
        xyz[2] = rzphi[1]
        return xyz

    def get_rz_of_psi(self, psi, y, x):
        self.psi_bisect = psi
        a = self.rmaxis
        b = self.rgefit[-1]
        r = bisect(self.f_rz_of_psi, a, b)
        z = self.zmaxis
        rz = np.array([r, z])
        return rz

    def f_rz_of_psi(self, r):
        z = self.zmaxis
        psi_rz = self.psi_rz_interp2d(r, z, grid = False)
        f = psi_rz - self.psi_bisect
        return f

    def segment_intersect(self, a, b, c, d):
        if not ( min(a[0],b[0]) <= max(c[0],d[0]) and min(c[1],d[1]) <= max(a[1],b[1]) and min(c[0], d[0]) <= max(a[0],b[0]) and min(a[1],b[1]) <= max(c[1],d[1]) ):
            intersection_point = np.array([0.0, 0.0])
            return False, intersection_point

        u = (c[0] - a[0]) * (b[1] - a[1]) - (b[0] - a[0]) * (c[1] - a[1])
        v = (d[0] - a[0]) * (b[1] - a[1]) - (b[0] - a[0]) * (d[1] - a[1])
        w = (a[0] - c[0]) * (d[1] - c[1]) - (d[0] - c[0]) * (a[1] - c[1])
        z = (b[0] - c[0]) * (d[1] - c[1]) - (d[0] - c[0]) * (b[1] - c[1])
        is_intersect = ( u * v <= 0.000000000000001 and w * z <= 0.000000000000001 )
        if not is_intersect:
            intersection_point = np.array([0.0, 0.0])
            return is_intersect, intersection_point
        else:
            intersection_point = np.array([0.0, 0.0])
            intersection_point[0] = ((b[0] - a[0]) * (c[0] - d[0]) * (c[1] - a[1]) - c[0] * (b[0] - a[0]) * (c[1] - d[1]) + a[0] * (b[1] - a[1]) * (c[0] - d[0])) / ((b[1] - a[1]) * (c[0] - d[0]) - (b[0] - a[0]) * (c[1] - d[1]))
            intersection_point[1] = ((b[1] - a[1]) * (c[1] - d[1]) * (c[0] - a[0]) - c[1] * (b[1] - a[1]) * (c[0] - d[0]) + a[1] * (b[0] - a[0]) * (c[1] - d[1])) / ((b[0] - a[0]) * (c[1] - d[1]) - (b[1] - a[1]) * (c[0] - d[0]))
            return is_intersect, intersection_point




    def plot_psi(self):
        ##inite the fig of matplotlib
        fig=plt.figure(figsize=(10,5))
        fig.subplots_adjust(top=0.9,bottom=0.1,wspace=0.5,hspace=0.55)

        ax0=fig.add_subplot(1,1,1)

        ct0 = ax0.contour(self.r, self.z, self.psirz, 15)
        
        line0 = ax0.plot(self.rlim, self.zlim, color="blue", linewidth=3)

        #plot LCFS
        line0 = ax0.plot(self.rbbbs, self.zbbbs, color="red", linewidth=2)

        #plot magnetic axis, and point of outer midplane LCFS
        ax0.scatter(self.rmaxis, self.zmaxis, color = "blue")
        ax0.scatter(self.rz_out_midplane_lcfs[0], self.rz_out_midplane_lcfs[1], color = "green")

        #plot firstwall and divertor
        '''
        try:
            line0 = ax0.plot(self.firstwall_inner[:, 0], self.firstwall_inner[:, 1], color="black", linewidth=3)
            line0 = ax0.plot(self.firstwall_outer[:, 0], self.firstwall_outer[:, 1], color="black", linewidth=3)
            line0 = ax0.plot(self.divertor[:, 0], self.divertor[:, 1], color="black", linewidth=3)
            line0 = ax0.plot(self.divertor[:, 0], -self.divertor[:, 1], color="black", linewidth=3)
        except:
            pass
        '''

        #plot magnetic field line from one point in rz coordinate
        point0 = np.array([2.04068068+ 0.01, -0.00681797])
        mfl0 = self.mflt_2d(point0)
        line0 = ax0.plot(mfl0[:, 0], mfl0[:, 1], color="green", linewidth=3)

        ax0.set_aspect('equal')
        fig.savefig("psi.png")
        fig.show()
        


    def plot_B(self):
        ##inite the fig of matplotlib
        fig=plt.figure(figsize=(10,6))
        fig.subplots_adjust(top=0.9,bottom=0.1,wspace=0.5,hspace=0.5)

        '''
        ax0=fig.add_subplot(2,2,1)

        #plot firstwall and divertor
        try:
            line0 = ax0.plot(self.firstwall_inner[:, 0], self.firstwall_inner[:, 1], color="black", linewidth=3)
            line0 = ax0.plot(self.firstwall_outer[:, 0], self.firstwall_outer[:, 1], color="black", linewidth=3)
            line0 = ax0.plot(self.divertor[:, 0], self.divertor[:, 1], color="black", linewidth=3)
            line0 = ax0.plot(self.divertor[:, 0], -self.divertor[:, 1], color="black", linewidth=3)
        except:
            pass

        #print(self.Br.shape, self.Br[62, 32], self.Bz[62, 32], self.rgefit[62], self.zgefit[32])
        position = np.array([1.9, 0.0])
        b = self.get_B_rz(position)
        print(position, b)
        #plot B
        q = ax0.quiver(self.r, self.z, self.Br, self.Bz, scale = 5)
        ax0.quiverkey(q, 2.2125, 0.0, 2, "B")

        ax0.scatter(1.9, 0.0)

        ax0.set_aspect('equal')
        '''

        ax0=fig.add_subplot(2,3,1)

        ax0.set_aspect('equal')
        #plot all wall
        line0 = ax0.plot(self.wall_all[:, 0], self.wall_all[:, 1], color="black", linewidth=3)
        #line0 = ax0.plot(self.wall_all_solps[:, 0], self.wall_all_solps[:, 1], color="black", linewidth=3)

        #plot LCFS
        line0 = ax0.plot(self.rbbbs, self.zbbbs, color="red", linewidth=2)

        #plot magnetic axis, and point of outer midplane LCFS
        ax0.scatter(self.rmaxis, self.zmaxis, color = "blue")
        ax0.scatter(self.rz_out_midplane_lcfs[0], self.rz_out_midplane_lcfs[1], color = "green")

        #plot x point
        ax0.scatter(self.rz_x_point[0], self.rz_x_point[1], color = "blue")
        ax0.set_ylabel('Z (m)')
        ax0.set_title("Wall")

        ax0=fig.add_subplot(2,3,2)
        cf0 = ax0.contourf(self.r, self.z, self.psirz)
        ax0.set_aspect('equal')
        fig.colorbar(cf0)
        ax0.set_title("Psi")


        ax0=fig.add_subplot(2,3,3)
        cf0 = ax0.contourf(self.r, self.z, self.Bmag)
        ax0.set_aspect('equal')
        fig.colorbar(cf0)
        ax0.set_title("Bmag")


        ax0=fig.add_subplot(2,3,4)
        cf0 = ax0.contourf(self.r, self.z, self.Br)
        ax0.set_aspect('equal')
        fig.colorbar(cf0)
        ax0.set_title("Br")

        ax0=fig.add_subplot(2,3,5)
        cf0 = ax0.contourf(self.r, self.z, self.Bz)
        ax0.set_aspect('equal')
        fig.colorbar(cf0)
        ax0.set_xlabel('R (m)')
        ax0.set_title("Bz")

        ax0=fig.add_subplot(2,3,6)
        cf0 = ax0.contourf(self.r, self.z, self.Bphi)
        ax0.set_aspect('equal')
        fig.colorbar(cf0)
        ax0.set_title("Bphi")

        fig.show()

    def read_pfcflux_magnetic(self, file_name, nr, nz, nphi):
        pfcflux_file = open(file_name)
        pfcflux_file.readline()
        pfcflux_file.readline()


        data = np.loadtxt(pfcflux_file)
        r = data[:, 0]
        r = r.reshape((nr, nz, nphi))
        r = r[0,:,:]
        print(r.shape)
        print(r[0:10, 0])

        z = data[:, 1]
        z = z.reshape((nr, nz, nphi))
        z = z[0,:,:]

        Br = data[:, 3]
        Br = Br.reshape((nr, nz, nphi))
        Br = Br[0,:,:]
        print("Br: ", Br.min(), Br.max())
        print(Br[0:10, 0])

        Bz = data[:, 4]
        Bz = Bz.reshape((nr, nz, nphi))
        Bz = Bz[0,:,:]

        Bphi = data[:, 5]
        Bphi = Bphi.reshape((nr, nz, nphi))
        Bphi = Bphi[0,:,:]

        psirz = data[:, 6]
        psirz = psirz.reshape((nr, nz, nphi))
        psirz = psirz[0,:,:]

        
        fig=plt.figure(figsize=(10,10))
        fig.subplots_adjust(top=0.9,bottom=0.1,wspace=0.5,hspace=0.55)

        ax0=fig.add_subplot(2,2,1)
        ax0.contourf(r, z, psirz)
        ax0.set_aspect('equal')

        ax0=fig.add_subplot(2,2,2)
        ax0.contourf(r, z, Br)
        ax0.set_aspect('equal')

        ax0=fig.add_subplot(2,2,3)
        ax0.contourf(r, z, Bz)
        ax0.set_aspect('equal')

        ax0=fig.add_subplot(2,2,4)
        ax0.contourf(r, z, Bphi)
        ax0.set_aspect('equal')

        fig.show()
        

    def plot_test(self):
        ##inite the fig of matplotlib
        fig=plt.figure(figsize=(10,5))
        fig.subplots_adjust(top=0.9,bottom=0.1,wspace=0.5,hspace=0.55)

        ax0=fig.add_subplot(1,1,1)


        #plot firstwall and divertor
        line0 = ax0.plot(self.firstwall_inner[:, 0], self.firstwall_inner[:, 1], color="black", linewidth=3)
        line0 = ax0.plot(self.firstwall_outer[:, 0], self.firstwall_outer[:, 1], color="black", linewidth=3)
        line0 = ax0.plot(self.divertor[:, 0], self.divertor[:, 1], color="black", linewidth=3)
        line0 = ax0.plot(self.divertor[:, 0], -self.divertor[:, 1], color="black", linewidth=3)

        val = np.zeros((self.nr, self.nz))
        for i in range(0, self.nr):
            for j in range(0, self.nz):
                val[i,j] = i * j

        print(self.r[0, 62], self.z[0, 62], val[0, 32])

        ax0.contourf(self.r, self.z, val)
        ax0.set_aspect('equal')
        #fig.show()

    def plot_profiles(self):
        ##inite the fig of matplotlib
        fig=plt.figure(figsize=(10,10))
        fig.subplots_adjust(top=0.9,bottom=0.1,wspace=0.5,hspace=0.55)

        ax0=fig.add_subplot(4,1,1)
        line0 = ax0.plot(self.psi_1d, self.pprime, color="black")

        ax0=fig.add_subplot(4,1,2)
        line0 = ax0.plot(self.psi_1d, self.pres, color="red")

        ax0=fig.add_subplot(4,1,3)
        line0 = ax0.plot(self.psi_1d, self.qpsi, color="blue")
        

        ax0=fig.add_subplot(4,1,4)
        line0 = ax0.plot(self.rgefit, self.fpol, color="green")
        

        fig.show()

    #plot magnetic field line 3d
    def plot_mfl3d(self):
        ##inite the fig of matplotlib
        fig=plt.figure(figsize=(10,5))
        fig.subplots_adjust(top=0.9,bottom=0.1,wspace=0.5,hspace=0.55)

        ax0=fig.gca(projection='3d')

        #plot magnetic field line from one point in rz coordinate
        point0 = np.array([2.40, 0.0, 0.0])
        mfl0 = self.mflt_3d(point0)
        line0 = ax0.plot(mfl0[:, 0], mfl0[:, 1], mfl0[:, 2])

        ax0.set_aspect('equal')

        max_range = np.array([mfl0[:, 0].max()-mfl0[:, 0].min(), mfl0[:, 1].max()-mfl0[:, 1].min(), mfl0[:, 2].max()-mfl0[:, 2].min()]).max() / 2.0
        mid_x = (mfl0[:, 0].max()+mfl0[:, 0].min()) * 0.5
        mid_y = (mfl0[:, 1].max()+mfl0[:, 1].min()) * 0.5
        mid_z = (mfl0[:, 2].max()+mfl0[:, 2].min()) * 0.5
        ax0.set_xlim(mid_x - max_range, mid_x + max_range)
        ax0.set_ylim(mid_y - max_range, mid_y + max_range)
        ax0.set_zlim(mid_z - max_range, mid_z + max_range)

        fig.show()

    def plot_sol2d(self):
        fig=plt.figure(figsize=(5,10))
        fig.subplots_adjust(top=0.9,bottom=0.1,wspace=0.5,hspace=0.55)

        ax0=fig.add_subplot(1,1,1)

        #plot all wall
        line0 = ax0.plot(self.wall_all[:, 0], self.wall_all[:, 1], color="black", linewidth=3)
        #line0 = ax0.plot(self.wall_all_solps[:, 0], self.wall_all_solps[:, 1], color="black", linewidth=3)

        #plot wall in gfile
        line0 = ax0.plot(self.rlim, self.zlim, color="blue", linewidth=3)

        #plot LCFS
        line0 = ax0.plot(self.rbbbs, self.zbbbs, color="red", linewidth=2)

        #plot magnetic axis, and point of outer midplane LCFS
        ax0.scatter(self.rmaxis, self.zmaxis, color = "blue")
        ax0.scatter(self.rz_out_midplane_lcfs[0], self.rz_out_midplane_lcfs[1], color = "green")

        #plot x point
        ax0.scatter(self.rz_x_point[0], self.rz_x_point[1], color = "blue")


        #plot magnetic field line from one point in rz coordinate, lamada: sol width
        lamada = 0.1
        point0 = np.array([self.rz_out_midplane_lcfs[0] + lamada/2.0, self.rz_out_midplane_lcfs[1]])
        mfl0 = self.mflt_2d_intersect_wall(point0)
        line0 = ax0.plot(mfl0[:, 0], mfl0[:, 1], color="green", linewidth=3)

    
        line0 = ax0.plot(self.inner_edge[:, 0], self.inner_edge[:, 1], color="red", linewidth=3)

        ax0.set_xlim(self.rmin, self.rmax)
        ax0.set_ylim(self.zmin, self.zmax)
        
        ax0.set_aspect('equal')
        fig.savefig("data_pfcflux/sol2d.png")
        plt.show()

    def plot_sol1d(self):
        fig=plt.figure(figsize=(5,10))
        fig.subplots_adjust(top=0.9,bottom=0.1,wspace=0.5,hspace=0.55)

        ax0=fig.add_subplot(1,1,1)

        #plot all wall
        line0 = ax0.plot(self.wall_all[:, 0], self.wall_all[:, 1], color="black", linewidth=3)
        #line0 = ax0.plot(self.wall_all_solps[:, 0], self.wall_all_solps[:, 1], color="black", linewidth=3)

        #plot wall in gfile
        line0 = ax0.plot(self.rlim, self.zlim, color="blue", linewidth=3)

        #plot LCFS
        line0 = ax0.plot(self.rbbbs, self.zbbbs, color="red", linewidth=2)

        #plot magnetic axis, and point of outer midplane LCFS
        ax0.scatter(self.rmaxis, self.zmaxis, color = "blue")
        ax0.scatter(self.rz_out_midplane_lcfs[0], self.rz_out_midplane_lcfs[1], color = "green")

        #plot x point
        ax0.scatter(self.rz_x_point[0], self.rz_x_point[1], color = "blue")



        '''
        #plot magnetic field line from one point in rz coordinate, lamada: sol width
        lamada = 0.02
        point0 = np.array([self.rz_out_midplane_lcfs[0] + lamada/2.0, self.rz_out_midplane_lcfs[1]])
        mfl0 = self.mflt_2d_intersect_wall(point0)
        #line0 = ax0.plot(mfl0[:, 0], mfl0[:, 1], color="green", linewidth=3)

        number_core_left = 0
        while mfl0[number_core_left, 1] > self.rz_x_point[1]:
            number_core_left = number_core_left + 1
        mfl0_core_left = np.zeros((number_core_left - 1, 2))
        length_core_left = 0.0
        for i in range(0, number_core_left - 1):
            mfl0_core_left[i, 0] = mfl0[i, 0]
            mfl0_core_left[i, 1] = mfl0[i, 1]
            if i > 0:
                dx = mfl0_core_left[i, 0] - mfl0_core_left[i-1, 0]
                dy = mfl0_core_left[i, 1] - mfl0_core_left[i-1, 1]
                length_core_left = length_core_left + math.sqrt(dx * dx + dy * dy)

        mfl0_divertor_left = np.zeros((mfl0.shape[0] - mfl0_core_left.shape[0] + 1, 2))
        length_divertor_left = 0.0
        ii = 0
        for i in range(number_core_left - 2, mfl0.shape[0]):
            mfl0_divertor_left[ii] = mfl0[i]
            if i > number_core_left - 2:
                dx = mfl0_divertor_left[ii, 0] - mfl0_divertor_left[ii-1, 0]
                dy = mfl0_divertor_left[ii, 1] - mfl0_divertor_left[ii-1, 1]
                length_divertor_left = length_divertor_left + math.sqrt(dx * dx + dy * dy)
            ii = ii + 1


        mfl0 = self.mflt_2d_intersect_wall(point0, reverse=True)
        #line0 = ax0.plot(mfl0[:, 0], mfl0[:, 1], color="green", linewidth=3)

        number_core_right = 0
        while mfl0[number_core_right, 1] > self.rz_x_point[1]:
            number_core_right = number_core_right + 1

        mfl0_core_right = np.zeros((number_core_right - 1, 2))
        length_core_right = 0.0
        for i in range(0, number_core_right - 1):
            mfl0_core_right[i, 0] = mfl0[i, 0]
            mfl0_core_right[i, 1] = mfl0[i, 1]
            if i > 0:
                dx = mfl0_core_right[i, 0] - mfl0_core_right[i-1, 0]
                dy = mfl0_core_right[i, 1] - mfl0_core_right[i-1, 1]
                length_core_right = length_core_right + math.sqrt(dx * dx + dy * dy)

        mfl0_divertor_right = np.zeros((mfl0.shape[0] - mfl0_core_right.shape[0] + 1, 2))
        length_divertor_right = 0.0
        ii = 0
        for i in range(number_core_right - 2, mfl0.shape[0]):
            mfl0_divertor_right[ii, 0] = mfl0[i, 0]
            mfl0_divertor_right[ii, 1] = mfl0[i, 1]
            if i > number_core_right - 2:
                dx = mfl0_divertor_right[ii, 0] - mfl0_divertor_right[ii-1, 0]
                dy = mfl0_divertor_right[ii, 1] - mfl0_divertor_right[ii-1, 1]
                length_divertor_right = length_divertor_right + math.sqrt(dx * dx + dy * dy)
            ii = ii + 1

        line0 = ax0.plot(mfl0_core_left[:, 0], mfl0_core_left[:, 1], color="red", linewidth=3)
        line0 = ax0.plot(mfl0_divertor_left[:, 0], mfl0_divertor_left[:, 1], color="green", linewidth=3)
        line0 = ax0.plot(mfl0_core_right[:, 0], mfl0_core_right[:, 1], color="red", linewidth=3)
        line0 = ax0.plot(mfl0_divertor_right[:, 0], mfl0_divertor_right[:, 1], color="green", linewidth=3)

        #line0 = ax0.plot(mfl1[:, 0], mfl1[:, 1], color="green", linewidth=3)

    
        #plot firstwall and divertor
        #line0 = ax0.plot(self.firstwall_inner[:, 0], self.firstwall_inner[:, 1], color="black", linewidth=3)
        #line0 = ax0.plot(self.firstwall_outer[:, 0], self.firstwall_outer[:, 1], color="black", linewidth=3)
        #line0 = ax0.plot(self.divertor[:, 0], self.divertor[:, 1], color="black", linewidth=3)
        #line0 = ax0.plot(self.divertor[:, 0], -self.divertor[:, 1], color="black", linewidth=3)


        #new plot one by one 
        #line0 = ax0.plot(self.divertor_upper_left[:, 0], self.divertor_upper_left[:, 1], color="black", linewidth=3)
        #line0 = ax0.plot(self.divertor_upper_middle[:, 0], self.divertor_upper_middle[:, 1], color="black", linewidth=3)
        #line0 = ax0.plot(self.divertor_upper_right[:, 0], self.divertor_upper_right[:, 1], color="black", linewidth=3)

        #line0 = ax0.plot(self.divertor_lower_left[:, 0], self.divertor_lower_left[:, 1], color="black", linewidth=3)
        #line0 = ax0.plot(self.divertor_lower_middle[:, 0], self.divertor_lower_middle[:, 1], color="black", linewidth=3)
        #line0 = ax0.plot(self.divertor_lower_right[:, 0], self.divertor_lower_right[:, 1], color="black", linewidth=3)


        print("======================== real structure ===================")
        length_core = length_core_left + length_core_right
        Psol = 0.216e6/2.0
        temperature = 100.0

        volume_sol = length_core * lamada * (2.0 * const.pi * self.rmaxis)

        dq = (Psol / volume_sol) / const.e
        dn = dq / temperature

        scale = 100.0
        length_simu = (length_divertor_left + length_core_left + length_core_right + length_divertor_right) / scale
        print("dn, dq: ", dn * scale, dq * scale)
        print("simulation length: ", length_simu)
        print("simu length1: ", length_divertor_left / scale)
        print("simu length2: ", (length_divertor_left + length_core_left + length_core_right) / scale)
        print("simu length3: ", (length_divertor_left + length_core_left + length_core_right + length_divertor_right) / scale)
        print("vedf position: ", length_simu * 0.01, ",", length_simu * 1.0 / 6.0, ",", length_simu * 2.0 / 6.0, ",", length_simu * 3.0 / 6.0, ",", length_simu * 4.0 / 6.0, ",", length_simu * 4.0 / 6.0, ",", length_simu * (1.0 - 0.01))

        print("======================== symmetry structure ===================")
        cell_length = 0.5e-5

        length_core = length_core_left + length_core_right
        length_divertor = length_divertor_left + length_divertor_right
        Psol = 0.216e6/2.0
        length_debye = 20.0e-3
        temperature = 100.0

        volume_sol = length_core * length_debye * (2.0 * const.pi * self.rmaxis)

        dq = (Psol / volume_sol) / const.e
        dn = dq / temperature

        scale = 100.0
        length_simu = (length_core + length_divertor) / scale
        print("dn, dq: ", dn * scale, dq * scale)
        print("simulation length: ", length_simu)
        print("source length1: ", 0.5 * length_divertor / scale)
        print("source length2: ", (0.5 * length_divertor + length_core) / scale)
        print("vedf position: ", 20.0 * cell_length, ",", length_simu * 1.0 / 6.0, ",", length_simu * 2.0 / 6.0, ",", length_simu * 3.0 / 6.0, ",", length_simu * 4.0 / 6.0, ",", length_simu * 4.0 / 6.0, ",", length_simu - 20.0 * cell_length)
        vedf_position = np.array([20.0 * cell_length, length_simu * 1.0 / 6.0, length_simu * 2.0 / 6.0, length_simu * 3.0 / 6.0, length_simu * 4.0 / 6.0, length_simu * 4.0 / 6.0, length_simu - 20.0 * cell_length])
        print("vedf index: ", vedf_position / cell_length)
        print("zoom region: ", 20.0*cell_length, length_simu - 20.0 * cell_length)
        '''

        line0 = ax0.plot(self.inner_edge[:, 0], self.inner_edge[:, 1], color="red", linewidth=3)
        
        ax0.set_aspect('equal')
        fig.savefig("efit/sol2d.png")
        plt.show()

    def output_for_pfcflux(self, angle_begin, angle_end, nphi):
        dphi = (angle_end - angle_begin) / (nphi - 1)        
        
        n_row = self.nr * self.nz * nphi

        data_pfcflux = np.zeros((n_row, 7))

        iter = 0
        for k in range(0, nphi):
            for j in range(0, self.nz):
                for i in range(0, self.nr):
                    data_pfcflux[iter, 0] = (i * self.drefit + self.rleft) * 1000.0
                    data_pfcflux[iter, 1] = (j * self.dzefit + self.zdown) * 1000.0
                    data_pfcflux[iter, 2] = angle_begin + k * dphi
                    data_pfcflux[iter, 3] = self.Br[i,j]
                    data_pfcflux[iter, 4] = self.Bz[i,j]
                    data_pfcflux[iter, 5] = self.Bphi[i,j]
                    data_pfcflux[iter, 6] = self.psirz[i,j] * (-2.0 * 3.1415926)
                    #data_pfcflux[iter, 6] = self.psirz_transpose[i,j]
                    iter = iter + 1

        #np.savetxt(self.gfile_name + "_pfcflux.txt3d", data_pfcflux)
        line1 = " Magnetic field is along phi-axis at phi=[-10 .. 10]"
        line2 = " R(mm)      Z(mm)       Phi(deg)  BR(T)    BZ(T)     BPhi(T) Psi(Wb/rad)"

        file_txt3d = open(self.gfile_name + "_pfcflux.txt3d", "w")
        
        file_txt3d.write(line1)
        file_txt3d.write('\n')
        file_txt3d.write(line2)
        file_txt3d.write('\n')

        for i in range(0, n_row):
            for j in range(0, 7):
                file_txt3d.write("{: .5f}".format(data_pfcflux[i,j]) + " ")
            file_txt3d.write('\n')


        B_omp = self.get_B_rz_3d(self.rz_out_midplane_lcfs)
        Bt = B_omp[2]
        Bp = math.sqrt(B_omp[0] * B_omp[0] + B_omp[1] * B_omp[1])
        print("--------------- pfcflux info -----------------------")
        print("Bt at omp: ", Bt)
        print("Bp at omp: ", Bp)

    def read_flux_wall(self, flux_wall_file):
        self.flux_wall = np.loadtxt(flux_wall_file)

    def cal_flux(self):

        self.phi0 = 152.0
        self.lamadaq = 0.05


        self.ref_line_start = self.rz_out_midplane_lcfs
        self.ref_line_end = np.array([self.rmax, self.rz_out_midplane_lcfs[1]])
        #self.ref_line_end = np.array([self.rmax, self.zmax])
        self.ref_psi_start = self.get_psi(self.ref_line_start)
        self.ref_psi_end = self.get_psi(self.ref_line_end)

        self.ref_n = 50
        self.ref_data = np.zeros((self.ref_n, 2))
        self.ref_data[0,0] = 0.0
        self.ref_data[0,1] = self.ref_psi_start
        self.ref_data[-1,0] = self.ref_line_end[0] - self.ref_line_start[0]
        self.ref_data[-1,1] = self.ref_psi_end       
        for i in range(1, self.ref_n - 1):
            r = self.ref_line_start[0] + i * (self.ref_line_end[0] - self.ref_line_start[0]) / self.ref_n
            z = self.ref_line_start[1] + i * (self.ref_line_end[1] - self.ref_line_start[1]) / self.ref_n
            distance = np.zeros(2)
            distance[0] = r - self.rz_out_midplane_lcfs[0]
            distance[1] = z - self.rz_out_midplane_lcfs[1]
            self.ref_data[i, 0] = r - self.ref_line_start[0]
            #self.ref_data[i, 0] = np.linalg.norm(distance)
            self.ref_data[i, 1] = self.get_psi(np.array([r,z]))

        self.flux_psi_r_interp1d = interpolate.interp1d(self.ref_data[:, 1], self.ref_data[:, 0])

        self.flux_wall_mesh_number = 10
        self.flux_wall_mesh = np.zeros(((self.flux_wall.shape[0]-1) * self.flux_wall_mesh_number, 2))
        self.flux_wall_mesh_norm = np.zeros((self.flux_wall.shape[0] * self.flux_wall_mesh_number, 3))
        i_node_global = 0
        for i in range(0, self.flux_wall.shape[0] - 1):
            point_start = self.flux_wall[i, :]
            point_end = self.flux_wall[i+1, :]
            for i_node in range(0, self.flux_wall_mesh_number):
                r = point_start[0] + (i_node + 0.5) * (point_end[0] - point_start[0]) / self.flux_wall_mesh_number
                z = point_start[1] + (i_node + 0.5) * (point_end[1] - point_start[1]) / self.flux_wall_mesh_number
                self.flux_wall_mesh[i_node_global, 0] = r
                self.flux_wall_mesh[i_node_global, 1] = z

                self.flux_wall_mesh_norm[i_node_global, 0] = point_end[1] - point_start[1]
                self.flux_wall_mesh_norm[i_node_global, 1] = point_start[0] - point_end[0]
                self.flux_wall_mesh_norm[i_node_global, 2] = 0.0
                mag = np.linalg.norm(self.flux_wall_mesh_norm[i_node_global,:])
                self.flux_wall_mesh_norm[i_node_global, :] = self.flux_wall_mesh_norm[i_node_global, :] / mag

                i_node_global = i_node_global + 1


        self.flux_data = np.zeros(self.flux_wall_mesh.shape[0])
        for i in range(0, self.flux_data.shape[0]):
            rz = self.flux_wall_mesh[i, :]
            psi = self.get_psi(rz)
            print(self.ref_psi_start, self.ref_psi_end, psi)
            if(psi < self.ref_psi_start or psi > self.ref_psi_end or rz[0] < self.rmin):
                self.flux_data[i] = 0.0
            else:

                r = self.flux_psi_r_interp1d(psi)
                print("r = ", r)

                normal = self.flux_wall_mesh_norm[i,:]
                print(self.rmin, self.rmax, rz[0])

                b = self.get_B_unit_rz_3d(rz)
                coef = abs(np.inner(normal, b))

                self.flux_data[i] = self.phi0 * math.exp(-r/self.lamadaq) * coef

        fig=plt.figure(figsize=(10,10))
        fig.subplots_adjust(top=0.9,bottom=0.1,wspace=0.5,hspace=0.55)

        ax0=fig.add_subplot(1,1,1)
        line0 = ax0.plot(self.flux_wall_mesh[:, 1], self.flux_data, color="black")
        #line0 = ax0.plot(self.flux_wall_mesh[:, 1], self.flux_wall_mesh[:, 0], color="black")

        fig.show()


            
if __name__ == '__main__':
    efit = Efit()
    efit.read_gfile("g037007.00501.hl2a")
    #efit.read_gfile("gfile.east")
    #efit.read_gfile("g022535.00830")  #sol1d 2a gfile
    #efit.read_gfile("g099996.00000case1_IP2500bp02")
    
    efit.read_firstwall_divertor("firstwall_inner_hl2a.txt", "firstwall_hl2a.txt", "divertor_hl2a.txt")
    efit.write_gfile("g037007.00501.hl2a.divertor")

    efit.init_interp()

    #efit.plot_psi()
    #efit.plot_profiles()
    #efit.plot_test()
    #efit.plot_mfl3d()
    efit.plot_sol2d()

    #pfcflux function
    #efit.plot_B()
    #efit.output_for_pfcflux(-10.0, 10.0, 41)
    #efit.read_pfcflux_magnetic("pfcflux_WEST_2014_400kA_Rext2p93m_Ip_scan_fixed_currents_rcv_CEDRES_eq_fields_for_PFCflux.txt3d", 41, 91, 91)
    #efit.read_pfcflux_magnetic("g099996.00000case1_IP2500bp02_pfcflux.txt3d", 41, 129, 129)


    #a = efit.get_B_rz(np.array([1.9, 0.0]))

    plt.show()
    #print(a)

